/*
  Copyright (c) 2010
  The Regents of the University of Michigan
  All Rights Reserved

  Permission is granted to use, copy, create derivative works, and
  redistribute this software and such derivative works for any purpose,
  so long as the name of the University of Michigan is not used in
  any advertising or publicity pertaining to the use or distribution
  of this software without specific, written prior authorization. If
  the above copyright notice or any other identification of the
  University of Michigan is included in any copy of any portion of
  this software, then the disclaimer below must also be included.

  This software is provided as is, without representation or warranty
  of any kind either express or implied, including without limitation
  the implied warranties of merchantability, fitness for a particular
  purpose, or noninfringement.  The Regents of the University of
  Michigan shall not be liable for any damages, including special,
  indirect, incidental, or consequential damages, with respect to any
  claim arising out of or in connection with the use of the software,
  even if it has been or is hereafter advised of the possibility of
  such damages.
*/
#include "lookup.h"

#include "bootstrap_nodes.h"
#include "clearview.h"
#include "clearview_internal.h"
#include "contact.h"
#include "id.h"
#include "network.h"
#include "utils.h"

#include <event.h>
#include <glib.h>
#include <stdio.h>
#include <string.h>

struct Lookup {
  // (can test for pending lookup by checking lookup_id)
  // All lookup containers containing contacts have ownership.
  char *lookup_id;
  struct ClearviewInstance *parent;
  int pending_finds_count;
  struct PendingFindNode pending_finds[LOOKUP_CONCURRENCY];

  /* Set of IDs of contacts that have replied, ordered by distance
     from this instance. (we don't actually use the values of lookups).
     Has ownership over contained IDs.
     WARNING: will autofree removed items. Remember to use steal. */
  GSequence *replied_contacts;

  /* Set of IDs of contacts we've queried. Has ownership over contained IDs.
     Values are all GINT_TO_POINTER(1).
     WARNING: will autofree removed items. Remember to use steal. */
  GHashTable *contacts_queried;

  /* Set of contacts we still need to query. Keys are IDs, values are
     the contacts. Must be kept sorted by distance from this instance.
     Contains contact_new()'d ContactAddresses. 
     WARNING: will autofree removed items. Remember to use steal. */
  GTree *contacts_to_query;
};


int clearview_lookup_is_done(struct Lookup *lookup) {
  return lookup->lookup_id == 0;
}


struct Lookup *lookup_new(void) {
  struct Lookup *ret = g_malloc(sizeof(struct Lookup));
  ret->lookup_id = NULL;
  // Parent pointer is so that pending finds can time out.
  int ii;
  for (ii = 0; ii != LOOKUP_CONCURRENCY; ++ii) {
    ret->pending_finds[ii].parent = ret;
    ret->pending_finds[ii].contact = NULL;
    ret->pending_finds[ii].timeout = NULL;
    ret->pending_finds[ii].tries_left = 0;
  }

  return ret;
}


void lookup_start(struct ClearviewInstance *cv, struct Lookup *lookup,
                  char *id) {
  int ii;
  static int bootstrap_next = 0;
  if (lookup->lookup_id) {
#ifndef LOG_ONLY_FATAL
    fprintf(stderr, "clearview_lookup called while old lookup still going!\n");
    return;
#endif
  }

  lookup->parent = cv;
  lookup->pending_finds_count = 0;
  lookup->lookup_id = id;

  lookup->replied_contacts = g_sequence_new((GDestroyNotify)clearview_id_free);
  lookup->contacts_to_query = g_tree_new_full(clearview_id_cmp_by_distance,
                                              lookup->lookup_id,
                                              NULL,
                                              (GDestroyNotify)contact_free);
  lookup->contacts_queried = g_hash_table_new_full(clearview_id_hash,
                                                   clearview_id_equal,
                                                   (GDestroyNotify)clearview_id_free,
                                                   NULL);

  for (ii = 0; ii != LOOKUP_CONCURRENCY;
       ii++, bootstrap_next = (bootstrap_next + 1) % (sizeof(bootstrap_nodes)
                                                      / sizeof(bootstrap_nodes[0]))) {
    struct ContactAddress *new_contact;
    // Never query ourselves.
    if (bootstrap_nodes[bootstrap_next].addr.s_addr == local_addr.s_addr) {
      continue;
    }
    new_contact = contact_new();
    contact_init(new_contact);
    *new_contact = bootstrap_nodes[bootstrap_next];
    g_assert(!new_contact->id);

    struct ContactAddress *existing_contact = g_tree_lookup(
      lookup->contacts_to_query, clearview_get_id(new_contact));
    if (existing_contact) {
      fprintf(stderr, "ID collision for nodes %s:%d and %s:%d",
              inet_ntoa(existing_contact->addr), existing_contact->port,
              inet_ntoa(new_contact->addr), new_contact->port);
    } else {
      g_tree_insert(lookup->contacts_to_query, clearview_get_id(new_contact),
                    new_contact);
    }
  }

  // To kick this off, send up to LOOKUP_CONCURRENCY (alpha) queries.
  for (ii = 0; ii != LOOKUP_CONCURRENCY; ++ii) {
    if (lookup_send_query(lookup) != 1)
      break;
  }
}


static void lookup_assert_valid(struct Lookup *lookup) {
  g_assert(lookup);
  g_assert(lookup->pending_finds_count >= 0);
  g_assert(lookup->pending_finds_count <= LOOKUP_CONCURRENCY);

  int ii;
  for (ii = 0; ii != LOOKUP_CONCURRENCY; ++ii) {
    g_assert(lookup->pending_finds[ii].parent == lookup);
    g_assert(!!lookup->pending_finds[ii].contact ==
             !!lookup->pending_finds[ii].timeout);
  }
}


void lookup_free(struct Lookup *victim) {
  lookup_assert_valid(victim);
  g_free(victim);
}


int lookup_should_continue(struct Lookup *lookup) {
  lookup_assert_valid(lookup);

  if (!lookup->lookup_id) {
#ifdef VERBOSE_LOG
    fputs("Lookup ended already!\n", stderr);
#endif
    return 0;
  }

  if (g_sequence_iter_is_end(
          g_sequence_get_begin_iter(lookup->replied_contacts))) {
    return 1;
  }

  if (!g_tree_nnodes(lookup->contacts_to_query)) {
    puts("Lookup starvation detected!");
    return 0;
  }

  char *closest_id = g_tree_get_min(lookup->contacts_to_query);
  char *farthest_responder = g_sequence_get(
    g_sequence_iter_prev(g_sequence_get_end_iter(lookup->replied_contacts)));
  if (clearview_id_cmp_by_distance(closest_id, farthest_responder,
                                   lookup->lookup_id) > 0 &&
      g_sequence_get_length(lookup->replied_contacts) >= LOOKUP_K) {
    lookup_assert_valid(lookup);
    return 0;
  } else {
    lookup_assert_valid(lookup);
    return 1;
  }
}


static void clearview_lookup_query_timeout(int fd, short ev_type, void *arg) {
  // It timed out, so we need to free its contact and remove it from the list.
  struct PendingFindNode *expired = arg;
  struct Lookup *lookup = expired->parent;
  struct ClearviewInstance *cv = lookup->parent;
  int ii;
  g_assert(expired);
  g_assert(cv);

  lookup_assert_valid(lookup);

  // The lookup finished.
  if (!lookup->lookup_id) {
    return;
  }

  logger_log_timeout(&cv->log, expired->contact);

#ifdef VERBOSE_LOG
  fprintf(stderr, "Query to %s:%hu timed out.\n",
          inet_ntoa(expired->contact->addr), expired->contact->port);
#endif

  if (expired->tries_left-- > 0) {
    // Retry!
    // TODO(swolchok): Merge this code with lookup_send_query,
    // which needs cleanup badly.
    struct timeval tv = {10, 0};
    if (clearview_send_find_node_request(cv, expired->contact,
                                         lookup->lookup_id) < 0) {
      goto CLEANUP;
    } else {
      struct event *victim = expired->timeout;
      expired->timeout = g_malloc(sizeof(*expired->timeout));
      evtimer_set(expired->timeout, clearview_lookup_query_timeout,
                  expired);
      evtimer_add(expired->timeout, &tv);

      evtimer_del(victim);
      g_free(victim);
      lookup_assert_valid(lookup);
      return;
    }
  }

CLEANUP:
  remove_pending_find(lookup, expired);
  for (ii = lookup->pending_finds_count; ii != LOOKUP_CONCURRENCY; ++ii) {
    if (lookup_send_query(lookup) != 1)
      break;
  }

  lookup_assert_valid(lookup);
}


static struct PendingFindNode *lookup_get_pending_find(struct Lookup *lookup) {
  lookup_assert_valid(lookup);
  g_return_val_if_fail(lookup->pending_finds_count < LOOKUP_CONCURRENCY, NULL);

  int ii;
  for (ii = 0; ii != LOOKUP_CONCURRENCY; ++ii) {
    if (!lookup->pending_finds[ii].timeout) {
      lookup->pending_finds_count++;
      return &lookup->pending_finds[ii];
    }
  }

  g_return_val_if_reached(NULL);
}


// If there are LOOKUP_CONCURRENCY (or more) queries running already, no-op.
// Also no-op if there are no contacts to query right now.
// Otherwise, start another find_node query.

// Return 1 if we started another query, 0 if we didn't because there
// were already too many queries running or no contacts left or the lookup should halt, and -1
// if we didn't because there are no queries running and no contacts
// left to query (i.e., time to stop).
int lookup_send_query(struct Lookup *lookup) {
  struct PendingFindNode *new_pending;

  lookup_assert_valid(lookup);

  if (lookup->pending_finds_count == LOOKUP_CONCURRENCY) {
    return 0;
  }

  if (!lookup->lookup_id) {
    return -1;
  }

  char *closest_id_to_query = g_tree_get_min(lookup->contacts_to_query);

  if (!closest_id_to_query) {
    return lookup->pending_finds_count ? 0 : -1;
  }


  // TODO(swolchok): check whether this test causes perf issues.
  if (!lookup_should_continue(lookup)) {
    return 0;
  }

  struct ContactAddress *to_query = g_tree_lookup(lookup->contacts_to_query,
                                                  closest_id_to_query);
  g_assert(to_query);
  g_assert(clearview_get_id(to_query) == closest_id_to_query);
  g_assert(to_query->addr.s_addr != local_addr.s_addr);
  g_tree_steal(lookup->contacts_to_query, closest_id_to_query);

  // Actually fire off the query first. If this doesn't work, we can bail out
  // and not make a mess to clean up.
  if (clearview_send_find_node_request(lookup->parent, to_query,
                                       lookup->lookup_id) < 0) {
    return 0;
  }

  g_hash_table_insert(lookup->contacts_queried,
                      clearview_id_copy(closest_id_to_query),
                      GINT_TO_POINTER(1));

  // Remember that we are waiting on this find.
  new_pending = lookup_get_pending_find(lookup);
  g_assert(!new_pending->contact);
  new_pending->contact = to_query;
  g_assert(!new_pending->timeout);
  new_pending->timeout = g_malloc(sizeof(*new_pending->timeout));
  new_pending->tries_left = MAX_RETRIES;

  evtimer_set(new_pending->timeout, clearview_lookup_query_timeout,
              new_pending);

  // Many Bothans died to bring us this timeout.
  // com/aelitiz/azureus/plugins/dht/impl/DHTPluginImpl.java;
  // passed to createUDP in milliseconds.
  struct timeval tv = {10, 0};

  evtimer_add(new_pending->timeout, &tv);
  lookup_assert_valid(lookup);
  return 1;
}


void lookup_cleanup(struct Lookup *lookup) {
  int ii;
  lookup_assert_valid(lookup);

  for (ii = 0; ii != LOOKUP_CONCURRENCY; ++ii) {
    if (lookup->pending_finds[ii].contact) {
      free_pending_find(&lookup->pending_finds[ii]);
    }
  }

  lookup->pending_finds_count = 0;

  lookup->lookup_id = NULL;

  g_sequence_free(lookup->replied_contacts);
  lookup->replied_contacts = NULL;

  g_hash_table_destroy(lookup->contacts_queried);
  lookup->contacts_queried = NULL;

  g_tree_destroy(lookup->contacts_to_query);
  lookup->contacts_to_query = NULL;
}


void remove_pending_find(struct Lookup *lookup,
                         struct PendingFindNode *victim) {

  lookup_assert_valid(lookup);
  g_return_if_fail(victim->contact && victim->timeout);
  int ii;
  // Find the victim RPC.
  for (ii = 0; ii != LOOKUP_CONCURRENCY; ++ii) {
    if (&lookup->pending_finds[ii] == victim) {
      break;
    }
  }
  // It had damn well better be in there.
  g_assert(ii != LOOKUP_CONCURRENCY);

  free_pending_find(victim);

  // Mark the victim RPC as deleted.
  memset(victim, 0, sizeof(struct PendingFindNode));
  victim->parent = lookup;

  lookup->pending_finds_count--;
  lookup_assert_valid(lookup);
}


void free_pending_find(struct PendingFindNode *victim) {
  g_assert(victim->contact);
  evtimer_del(victim->timeout);
  contact_free(victim->contact);
  victim->contact = NULL;
  g_free(victim->timeout);
  victim->timeout = NULL;
}


struct BootstrapData {
  struct ClearviewInstance *cv;
  struct event this_event;
};


static void bootstrap_wrapper(int fd, short ev_type, void *arg) {
  struct BootstrapData *data = arg;
  struct ClearviewInstance *cv = data->cv;

  evtimer_del(&data->this_event);
  g_slice_free(struct BootstrapData, data);
  clearview_join_network(cv);
}


void lookup_handle_find_reply(struct Lookup *lookup,
                              struct ReplyFindNodePacket *reply,
                              struct ContactAddress *contact) {
  // TODO(swolchok): refcount contacts in the routing tables and lookup
  //                 structures.
  int ii, num_replied;
  char *to_insert;
  struct ClearviewInstance *cv = lookup->parent;

  lookup_assert_valid(lookup);

  // Is the lookup over already?
  if (clearview_lookup_is_done(lookup)) {
    goto END;
  }

  // Contact is not heap-allocated, so we have ownership over this ID.
  g_assert(!contact->id);  // Make sure we have ownership.
  to_insert = clearview_get_id(contact);
  g_sequence_insert_sorted(lookup->replied_contacts, to_insert,
                           clearview_id_cmp_by_distance, lookup->lookup_id);
  num_replied = g_sequence_get_length(lookup->replied_contacts);

  // Only need the K closest.
  if (num_replied > LOOKUP_K) {
    GSequenceIter *victim_iter = g_sequence_iter_prev(
        g_sequence_get_end_iter(lookup->replied_contacts));
    // Container has ownership over contained IDs; DO need to free it.
    // However, we set up clearview_id_free as the destroy function, so
    // we don't have to call it explicitly here.
    g_assert(num_replied == LOOKUP_K+1);
    g_sequence_remove(victim_iter);
    num_replied--;
  }

  // Remove this contact from the list of pending find_nodes.
  for (ii = 0; ii != LOOKUP_CONCURRENCY; ++ii) {
    struct ContactAddress *candidate = lookup->pending_finds[ii].contact;
    if (!candidate) {
      // Candidate marked as deleted.
      continue;
    }
    if (candidate->addr.s_addr == contact->addr.s_addr &&
        candidate->port == contact->port) {
      remove_pending_find(lookup, &lookup->pending_finds[ii]);
      break;
    }
  }

  for (ii = 0; ii != reply->contacts_count; ++ii) {
    // Never query ourselves.
    if (reply->contacts[ii]->addr.s_addr == local_addr.s_addr) {
      contact_free(reply->contacts[ii]);
      continue;
    }

    logger_log_contact(&cv->log, reply->contacts[ii]);
    char *id = clearview_get_id(reply->contacts[ii]);
    // Don't add contacts we've already queried or already know to query.
    if (g_hash_table_lookup(lookup->contacts_queried, id) ||
        g_tree_lookup(lookup->contacts_to_query, id)) {
      contact_free(reply->contacts[ii]);
    } else {
      g_tree_insert(lookup->contacts_to_query, id, reply->contacts[ii]);
    }
  }

  if (!lookup_should_continue(lookup)) {
    struct timeval one_hour = {60 * 60, 0};
    struct BootstrapData *data = g_slice_new(struct BootstrapData);
    time_t the_time = time(NULL);

    printf("Lookup for port %hu finished at %s", cv->port,
           ctime(&the_time));
    fflush(stdout);

    data->cv = lookup->parent;
    evtimer_set(&data->this_event, bootstrap_wrapper, data);
    evtimer_add(&data->this_event, &one_hour);
    lookup_cleanup(lookup);
    goto END;
  }

  lookup_assert_valid(lookup);
  for (ii = lookup->pending_finds_count; ii != LOOKUP_CONCURRENCY; ++ii) {
    if (lookup_send_query(lookup) != 1) {
      break;
    }
  }

END:
  g_free(reply->contacts);
  lookup_assert_valid(lookup);
}
